import sys

import pandas as pd
from torch.utils.data import Dataset, DataLoader
import joblib
import os
import torch
import numpy as np
import random
import matplotlib.pyplot as plt
from orbitP.script import config

class orbitPSULTDataset(Dataset):
    """
    测试数据集：使用前一天数据预测后一天，连续进行X天预测
    """
    def __init__(self, obsData,prdData,stampObs,stampPrd, training_length, predicting_length, forecast_window, days):
        self.obsData = obsData
        self.prdData = prdData
        self.stampObs = stampObs
        self.stampPrd = stampPrd
        self.T = training_length
        self.P = predicting_length
        self.S = forecast_window
        self.days = int(days)
        self.obsIdx = [i * config.training_length for i in range(self.days)]
        self.prdIdx = [i * config.predicting_length for i in range(self.days)]

    def __len__(self):
        return self.days

    def __getitem__(self, idx):
        startObs = self.obsIdx[idx]
        startPrd = self.prdIdx[idx]
        orbitData_pre = torch.tensor(self.obsData[startObs:startObs+self.T, :])
        orbitData_suf = torch.tensor(self.prdData[startPrd:startPrd+self.P, :])
        stampData_pre = torch.tensor(self.stampObs[startObs:startObs+self.T, :])
        stampData_suf = torch.tensor(self.stampPrd[startPrd:startPrd+self.P, :])

        return orbitData_pre, orbitData_suf, stampData_pre, stampData_suf